package com.example.demo.util;

import javax.sql.DataSource;
import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.sql.*;
import java.util.List;
import java.util.Scanner;

import org.json.JSONException;
import org.json.JSONObject;
import org.springframework.stereotype.Component;

@Component
public class AIGenerateImageUtil {
    private DataSource dataSource; // 假设已经注入了数据源

    public Blob callAIGenerateImageLogic(List<String> selectedIngredients) {
        String apiKey = "";  // 从控制台获取
        String apiSecret = "";  // 从控制台获取
        String query = "";  // 图片主题
        String style = "";  // 图片风格
        String requestUrl = "https://pretrain.aminer.cn/api/v2/cogview";

        // 构造请求参数
        String jsonInputString = "{\"query\": \"" + query + "\", \"style\": \"" + style + "\", \"apikey\": \"" + apiKey + "\", \"apisecret\": \"" + apiSecret + "\"}";

        // 发送POST请求
        try {
            byte[] generatedImageData = sendPostRequest(requestUrl, jsonInputString);

            try (Connection connection = dataSource.getConnection()) {
                String insertSql = "INSERT INTO generated_images (image_data) VALUES (?)";
                try (PreparedStatement statement = connection.prepareStatement(insertSql, Statement.RETURN_GENERATED_KEYS)) {
                    ByteArrayInputStream inputStream = new ByteArrayInputStream(generatedImageData);
                    statement.setBinaryStream(1, inputStream, generatedImageData.length);
                    int affectedRows = statement.executeUpdate();
                    if (affectedRows == 0) {
                        throw new SQLException("Creating image failed, no rows affected.");
                    }
                    try (ResultSet generatedKeys = statement.getGeneratedKeys()) {
                        if (generatedKeys.next()) {
                            Blob generatedImageBlob = generatedKeys.getBlob(1);
                            return generatedImageBlob;
                        } else {
                            throw new SQLException("Creating image failed, no ID obtained.");
                        }
                    }
                }
            } catch (SQLException e) {
                // 处理异常
                e.printStackTrace();
                return null; // 返回空值或者抛出异常，视情况而定
            }
        } catch (IOException | InterruptedException e) {
            e.printStackTrace();
            return null;
        }
    }

    private byte[] sendPostRequest(String requestUrl, String jsonInputString) throws IOException, InterruptedException {
        URL url = new URL(requestUrl);
        HttpURLConnection connection = (HttpURLConnection) url.openConnection();
        connection.setRequestMethod("POST");
        connection.setRequestProperty("Content-Type", "application/json");
        connection.setDoOutput(true);

        try (OutputStream os = connection.getOutputStream()) {
            byte[] input = jsonInputString.getBytes("utf-8");
            os.write(input, 0, input.length);
        }

        // 处理响应
        if (connection.getResponseCode() == HttpURLConnection.HTTP_OK) {
            try (Scanner scanner = new Scanner(connection.getInputStream(), "utf-8")) {
                String responseBody = scanner.useDelimiter("\\A").next();
                System.out.println(responseBody);
                String taskId = parseTaskId(responseBody);
                if (taskId != null) {
                    String statusRequestUrl = "https://pretrain.aminer.cn/api/v2/status?task_id=" + taskId;
                    checkStatus(statusRequestUrl);
                    return getGeneratedImageData(taskId); // 获取生成的图片数据
                }
            }
        } else {
            System.out.println("Error: " + connection.getResponseCode() + " " + connection.getResponseMessage());
        }

        return null;
    }

    private String parseTaskId(String responseBody) {
        try {
            JSONObject jsonObject = new JSONObject(responseBody);
            return jsonObject.getJSONObject("result").getString("task_id");
        } catch (JSONException e) {
            e.printStackTrace();
            return null;
        }
    }

    private void checkStatus(String statusRequestUrl) throws IOException, InterruptedException {
        // 发送GET请求检查任务状态
        while (true) {
            URL url = new URL(statusRequestUrl);
            HttpURLConnection connection = (HttpURLConnection) url.openConnection();
            connection.setRequestMethod("GET");

            if (connection.getResponseCode() == HttpURLConnection.HTTP_OK) {
                try (Scanner scanner = new Scanner(connection.getInputStream(), "utf-8")) {
                    String responseBody = scanner.useDelimiter("\\A").next();
                    System.out.println(responseBody);
                    if (isTaskComplete(responseBody)) {
                        break;
                    } else {
                        Thread.sleep(10000);  // 等待10秒后再次检查任务状态
                    }
                }
            } else {
                System.out.println("Error: " + connection.getResponseCode() + " " + connection.getResponseMessage());
                break;
            }
        }
    }

    private boolean isTaskComplete(String responseBody) {
        try {
            JSONObject jsonObject = new JSONObject(responseBody);
            String status = jsonObject.getJSONObject("result").getString("status");
            return "completed".equals(status);
        } catch (JSONException e) {
            e.printStackTrace();
            return false;
        }
    }

    private byte[] getGeneratedImageData(String taskId) {
        // 使用任务ID获取生成的图片数据
        String requestUrl = "https://pretrain.aminer.cn/api/v2/result?task_id=" + taskId;

        // 发送GET请求获取生成的图片数据
        try {
            URL url = new URL(requestUrl);
            HttpURLConnection connection = (HttpURLConnection) url.openConnection();
            connection.setRequestMethod("GET");

            if (connection.getResponseCode() == HttpURLConnection.HTTP_OK) {
                try (InputStream inputStream = connection.getInputStream()) {
                    ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
                    byte[] buffer = new byte[1024];
                    int bytesRead;
                    while ((bytesRead = inputStream.read(buffer)) != -1) {
                        outputStream.write(buffer, 0, bytesRead);
                    }
                    return outputStream.toByteArray();
                }
            } else {
                System.out.println("Error: " + connection.getResponseCode() + " " + connection.getResponseMessage());
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        return new byte[0]; // 返回空字节数组
    }
}